#!/usr/bin/env python
# encoding: utf-8
'''
@author: Excelsiorly
@license: (C) Copyright 2022, All Rights Reserved.
@contact: excelsiorly@qq.com
@file: 498. 对角线遍历.py
@time: 2022/3/7 13:17
@desc: https://leetcode-cn.com/problems/diagonal-traverse/
> 给你一个大小为 m x n 的矩阵 mat ，请以对角线遍历的顺序，用一个数组返回这个矩阵中的所有元素。

1. 找规律, x+y==趟数i(从0开始），奇数趟是x从i开始递减，y从i-x开始递增，一直到越界，偶数趟y从趟数i开始，x从i-y开始，直到越界
2. 见题解：https://leetcode-cn.com/problems/diagonal-traverse/solution/dui-jiao-xian-bian-li-fen-xi-ti-mu-zhao-zhun-gui-l/
'''


class Solution:
    def findDiagonalOrder(self, mat: List[List[int]]) -> List[int]:
        m, n = len(mat), len(mat[0])
        nums = []
        i = 0
        while i < m + n:
            # 1 3 5 ... 趟
            x1 = i if i < m else m - 1
            y1 = i - x1
            while x1 >= 0 and y1 < n:
                nums.append(mat[x1][y1])
                x1 -= 1
                y1 += 1
            i += 1

            if i >= m + n: break

            # 2 4 6 ... 趟
            y2 = i if i < n else n - 1
            x2 = i - y2
            while y2 >= 0 and x2 < m:
                nums.append(mat[x2][y2])
                x2 += 1
                y2 -= 1
            i += 1
        return nums